# %%
import datetime
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from docx.shared import Cm
from docxtpl import DocxTemplate, InlineImage
from scipy.optimize import fsolve

SCRIPT_TO_INSERT_PATH = "sympy2subdoc.py"

# Исходные данные
P_0 = 50 * 10**5  # Па
P_H = 10**5  # Па
T_0 = 550  # К
M_A = 3.5
D_A = 0.15  # м
K = 1.25
R = 313  # Дж/(кг*К)
D_0 = D_A
L_0 = 0.2 * D_A  # В предыдущих обозначениях просто l(иначе ruff ругается)

# Снизу газодинамические функции


# Сверху газодинамические функции
# %%
# Снизу функции отображения графиков
y_labels = {
    "давление": r"$ P(\widetilde{x}), Па $",
    "температура": r"$ T(\widetilde{x}), К $",
    "мах": r"$ M(\widetilde{x}) $",
    "плотность": r"$ \rho(\widetilde{x}), кг/м^3 $",
    "лямбда": r"$\lambda(\widetilde{x})$",
    "скорость": r"$ V(\widetilde{x}), a(\widetilde{x}), м/с $",
}

font_size = 16
img_path = r"dz/isentropic_steady_flow/"


def graph(
    x: np.ndarray,
    *functions_y: np.ndarray,
    y_label: str = None,
    legend: list = None,
):
    """Отрисовка графиков

    Parameters
    ----------
    x : np.ndarray
        массив чисел координаты x

    functions_y: np.ndarray
        один или несколько массивов данных, которые необходимо отложить по оси y, можно
        просто через запятую их вписать в функцию

    y_label = : str, optional
        by default None

        выбрать подпись графика, возможные варианты: ['давление','температура',
        'мах','плотность','лямбда','скорость']

    legend = : list, optional
        by default None

        массив из строк с подписями, которые должны быть отражены в легенде графика,
        например на графике скорости надо подписать legend=['V','a','a*'], а не как имя
        переменной, которой график обозначен. Размер массива должен совпадать с
        количеством выводимых графиков

    Пример
        graph(x, y, z, y_label = 'скорость', legend = ['y', 'z'])
    """
    if y_label is not None:
        y_label = y_label.lower()
    functions = [function for function in functions_y]
    fig, ax = plt.subplots(figsize=(13, 8))
    if legend is None:
        for i in range(len(functions)):
            ax.plot(x, functions[i])
    else:
        for i in range(len(functions)):
            ax.plot(x, functions[i], label=legend[i])
        ax.legend(loc="best", fontsize=font_size)
    ax.set_xlabel(r"$\widetilde{x}$", fontsize=font_size)
    ax.set_ylabel(
        y_labels.get(y_label),
        fontsize=font_size,
        rotation="vertical",
        loc="center",
    )
    ax.ticklabel_format(useMathText=True)
    ax.yaxis.get_offset_text().set_fontsize(font_size)
    ax.tick_params(axis="both", labelsize=font_size)
    ax.set_xlim(x[0], x[-1])
    ax.set_ylim(bottom=0)
    plt.grid(which="major", axis="both")
    global img_number
    if y_label is not None:
        plt.savefig(img_path + y_label + ".png")
    else:
        plt.savefig(img_path + "форма_сопла" + ".png")


# Сверху функции отображения графиков
# %%
# Снизу расчет параметров


# Газодинамические функции
def m_to_lambda(M: np.ndarray) -> np.ndarray:
    """Перевод числа маха в относительную скорость

    Parameters
    ----------
    M : np.ndarray
        Число Маха, б/р

    Returns
    -------
    l : np.ndarray
        Относительная скорость, б/р
    """
    L = np.sqrt(((K + 1) / 2 * M**2) / (1 + (K - 1) / 2 * M**2))
    return L


def pi(L: np.ndarray) -> np.ndarray:
    """Газодинамическая функция Пи

    Parameters
    ----------
    l : np.ndarray
        Относительная скорость, б/р

    Returns
    -------
    Pi : np.ndarray
        Значение газодинамической функции Пи
    """
    Pi = (1 - L**2 * (K - 1) / (K + 1)) ** (K / (K - 1))
    return Pi


def eps(L: np.ndarray) -> np.ndarray:
    """Газодинамическая функция Эпсилон

    Parameters
    ----------
    l : np.ndarray
        Относительная скорость, б/р

    Returns
    -------
    Eps : np.ndarray
        Значение газодинамической функции Эпсилон
    """
    Eps = (1 - L**2 * (K - 1) / (K + 1)) ** (1 / (K - 1))
    return Eps


def tau(L: np.ndarray) -> np.ndarray:
    """Газодинамическая функция Тау

    Parameters
    ----------
    L : np.ndarray
        Относительная скорость, б/р

    Returns
    -------
    Tau : np.ndarray
        Значение газодинамической функции Тау
    """
    Tau = 1 - L**2 * (K - 1) / (K + 1)
    return Tau


def q(L: np.ndarray) -> np.ndarray:
    """Газодинамическая функция q

    Parameters
    ----------
    L : np.ndarray
        Относительная скорость, б/р

    Returns
    -------
    Q : np.ndarray
        Значение газодинамической функции Q
    """
    Q = (
        L
        * ((K + 1) / 2) ** (1 / (K - 1))
        * (1 - (K - 1) / (K + 1) * L**2) ** (1 / (K - 1))
    )
    return Q


def lambda_to_a(a0: float, akr: float, L: np.ndarray) -> np.ndarray:
    """
    Перевод относительной скорости в скорость звука

    Parameters
    ----------
    L : np.ndarray
        Относительная скорость, б/р

    a0 - скорость звука в сечении торможения, м/с

    akr - критическая скорость звука, м/с

    Returns
    -------
    A : np.ndarray
        Величина скорости звука в м/с
    """

    A = np.sqrt(0.5 * (K - 1) * (a0**2 * (2 / (K - 1)) - (akr * L) ** 2))

    return A


# Вычисления
n1, n2 = 2, 5  # число секций
x_r = np.linspace(0, (n1 + n2) * L_0, num=1000)
x_br = x_r / D_A
l_a = m_to_lambda(M_A)
a_0 = np.sqrt(K * R * T_0)  # скорсоть звука в емкости, м/с
ro_0 = P_0 / (R * T_0)  # плотность воздуха в емкости
p_a = P_0 * pi(l_a)  # давление на срезе сопла, Па
ro_a = ro_0 * eps(l_a)  # плотность на срезе сопла, кг/м**3
T_a = T_0 * tau(l_a)  # температура на срезе сопла, К
a_kr = a_0 * np.sqrt(2 / (K + 1))  # скорость звука в критике, м/с
ro_kr = ro_0 * (2 / (K + 1)) ** (1 / (K - 1))  # плотность в критике, кг/м**3
V_a = l_a * a_kr  # скорость на срезе сопла, м/с
a_a = V_a / M_A  # скорость звука на срезе сопла, м/с
d_kr = D_A * np.sqrt(q(l_a))  # диаметр критического сечения, м

# диаметр сопла в сечении, м
d_x_br = np.where(
    x_br < n1 * L_0 / D_A,
    D_0 + (d_kr - D_0) / (n1 * L_0) * D_A * x_br,
    d_kr + (D_A - d_kr) / (n2 * L_0) * D_A * (x_br - n1 * L_0 / D_A),
)

m_rash = ro_kr * a_kr * np.pi * d_kr**2 / 4  # расход газа, кг/с
S_d = np.pi * d_x_br**2 / 4  # площадь сопла в сечении, м**2

# вектор начала решения
lambda_zero = np.zeros(len(d_x_br)) + np.where(x_br < n1 * L_0 / D_A, 0, 2)
lambda_x = fsolve(lambda L: q(L) - d_kr**2 / d_x_br**2, lambda_zero)

# Сверху расчет параметров
# graph(x, y, z, y_label = 'скорость', legend = ['y', 'z'])
# Построение графиков

# Геометрия сопла
nozzle_geom = graph(x_br, d_x_br)

# график функции лямбда
lam_graph = graph(x_br, lambda_x, y_label="лямбда")

# график скорости и скорости звука, на котором дополнительно выводится критика
vel_and_a_graph = graph(
    x_br,
    lambda_x * a_kr,
    lambda_to_a(a_0, a_kr, lambda_x),
    np.full(len(x_br), a_kr),
    y_label="скорость",
    legend=[r"$V(\widetilde{x})$", r"$a(\widetilde{x})$", "a*"],
)

# график функции Давления
pressure_graph = graph(x_br, P_0 * pi(lambda_x), y_label="давление")

# график функции плотности
density_graph = graph(x_br, ro_0 * eps(lambda_x), y_label="плотность")

# график числа маха
mach_graph = graph(
    x_br,
    lambda_x * np.sqrt(1 / ((K + 1) / 2 - lambda_x**2 * (K - 1) / 2)),
    y_label="мах",
)

# график температуры
temp_graph = graph(x_br, T_0 * tau(lambda_x), y_label="температура")
# посмотреть необходимость использования легенды графиков

# Снизу вывод в файл
OUTPUT_FOLDER = Path("dz/isentropic_steady_flow")
DOCX_TEMPLATE_PATH = OUTPUT_FOLDER.joinpath("template.docx")
REPORT_OUTPUT_PATH = OUTPUT_FOLDER.joinpath("generated_report.docx")
TEMPORARY_IMAGE_PATH_NOZZLE_SHAPE = OUTPUT_FOLDER.joinpath("форма_сопла.png")
TEMPORARY_IMAGE_PATH_LAMBDA = OUTPUT_FOLDER.joinpath("лямбда.png")
TEMPORARY_IMAGE_PATH_PRESSURE = OUTPUT_FOLDER.joinpath("давление.png")
TEMPORARY_IMAGE_PATH_DENSITY = OUTPUT_FOLDER.joinpath("плотность.png")
TEMPORARY_IMAGE_PATH_TEMPERATURE = OUTPUT_FOLDER.joinpath("температура.png")
TEMPORARY_IMAGE_PATH_SPEED = OUTPUT_FOLDER.joinpath("скорость.png")
TEMPORARY_IMAGE_PATH_MAX = OUTPUT_FOLDER.joinpath("мах.png")

docx_template = DocxTemplate(DOCX_TEMPLATE_PATH)

initial_data = {
    "parameters": [
        "P\u2080",
        "T\u2080",
        "M\u2090",
        "d\u2090",
        "k",
        "R",
        "d\u2080",
        "l\u2080",
    ],
    "values": [
        f"{P_0/10**5} * 10\u2075 Па",
        f"{T_0} К",
        f"{M_A}",
        f"{D_A} м",
        f"{K}",
        f"{R} Дж/(кг*К)",
        "d\u2090",
        "0.2 * d\u2090",
    ],
}

table_initial_parameters = []
x = []
y = []

for i in range(len(list(initial_data["parameters"]))):
    table_initial_parameters.append(
        {
            "parameter": (initial_data["parameters"][i]),
            "value": initial_data["values"][i],
        }
    )
    x.append(initial_data["parameters"][i])
    y.append(initial_data["values"][i])

image_nozzle_shape_sizes = InlineImage(
    docx_template, str(TEMPORARY_IMAGE_PATH_NOZZLE_SHAPE), Cm(12)
)

image_lambda = InlineImage(docx_template, str(TEMPORARY_IMAGE_PATH_LAMBDA), Cm(12))

table_lambda_values = []
h = []
la = []
step = 1000

for i in range(0, step, 143):
    table_lambda_values.append(
        {
            "index": f"{x_br[i]:.1f}",
            "value": f"{lambda_x[i]:.3f}",
        }
    )
    h.append(i)
    la.append(lambda_x[i])

table_lambda_values.append(
    {
        "index": f"{x_br[-1]:.1f}",
        "value": f"{lambda_x[-1]:.3f}",
    }
)  # вводим в таблицу последнее значение


image_pressure = InlineImage(docx_template, str(TEMPORARY_IMAGE_PATH_PRESSURE), Cm(12))
image_density = InlineImage(docx_template, str(TEMPORARY_IMAGE_PATH_DENSITY), Cm(12))
image_temperature = InlineImage(
    docx_template, str(TEMPORARY_IMAGE_PATH_TEMPERATURE), Cm(12)
)
image_speed = InlineImage(docx_template, str(TEMPORARY_IMAGE_PATH_SPEED), Cm(12))
image_max = InlineImage(docx_template, str(TEMPORARY_IMAGE_PATH_MAX), Cm(12))

n_parameter = f"n = {(p_a / P_H):.3f}"

received_data = {
    "parameters": [
        "a\u2080",
        "\u03C1\u2080",
        "\u03BB\u2090",
        "P\u2090",
        "\u03C1\u2090",
        "T\u2090",
        "a*",
        "V\u2090",
        "a\u2090",
        "d*",
        "m",
    ],
    "values": [
        round(a_0, 3),
        round(ro_0, 3),
        round(l_a, 3),
        round(p_a, 3),
        round(ro_a, 3),
        round(T_a, 3),
        round(a_kr, 3),
        round(V_a, 3),
        round(a_a, 3),
        round(d_kr, 3),
        round(m_rash, 3),
    ],
    "dimension": [
        "м/с",
        "кг/м\u00b3",
        " - ",
        "Па",
        "кг/м\u00b3",
        "K",
        "м/с",
        "м/с",
        "м/с",
        "м",
        "кг/с",
    ],
}

table_received_values = []
pp = []
vv = []
dd = []

for i in range(len(list(received_data["parameters"]))):
    table_received_values.append(
        {
            "parameter": (received_data["parameters"][i]),
            "value": received_data["values"][i],
            "dimension": received_data["dimension"][i],
        }
    )
    pp.append(received_data["parameters"][i])
    vv.append(received_data["values"][i])
    dd.append(received_data["dimension"][i])

table_functions = []
z = []
p = []
d = []
t = []
s = []
ss = []
m = []
step = 1000

for i in range(0, step, 143):
    num_pressure = ((P_0 * pi(lambda_x))[i]) / 10**6
    num_density = (ro_0 * eps(lambda_x))[i]
    num_temperature = (T_0 * tau(lambda_x))[i]
    num_speed = lambda_x[i] * a_kr
    num_speedsound = lambda_to_a(a_0, a_kr, lambda_x)[i]
    num_max = lambda_x[i] * np.sqrt(2 / (K + 1 - lambda_x[i] ** 2 * (K - 1)))
    table_functions.append(
        {
            "index": f"{x_br[i]:.1f}",
            "pressure": f"{num_pressure:.3f}",
            "density": f"{num_density:.3f}",
            "temperature": f"{num_temperature:.3f}",
            "speed": f"{num_speed:.3f}",
            "speedsound": f"{num_speedsound:.3f}",
            "max": f"{num_max:.3f}",
        }
    )
    z.append(i)
    p.append(num_pressure)
    d.append(num_density)
    t.append(num_temperature)
    s.append(num_speed)
    ss.append(num_speedsound)
    m.append(num_max)

formula_max = lambda_x[-1] * np.sqrt(2 / (K + 1 - lambda_x[-1] ** 2 * (K - 1)))

table_functions.append(
    {
        "index": f"{x_br[-1]:.1f}",
        "pressure": f"{((P_0 * pi(lambda_x))[-1])/10**6:.3f}",
        "density": f"{(ro_0 * eps(lambda_x))[-1]:.3f}",
        "temperature": f"{(T_0 * tau(lambda_x))[-1]:.3f}",
        "speed": f"{(lambda_x[-1] * a_kr):.3f}",
        "speedsound": f"{(lambda_to_a(a_0, a_kr, lambda_x)[-1]):.3f}",
        "max": f"{formula_max:.3f}",
    }
)  # вводим в таблицу последнее значение

context = {
    "year": datetime.datetime.now().strftime("%Y"),
    "image_nozzle_shape_sizes": image_nozzle_shape_sizes,
    "table_initial_parameters": table_initial_parameters,
    "image_lambda": image_lambda,
    "table_lambda_values": table_lambda_values,
    "image_pressure": image_pressure,
    "image_density": image_density,
    "image_temperature": image_temperature,
    "image_speed": image_speed,
    "image_max": image_max,
    "n_parameter": n_parameter,
    "table_received_values": table_received_values,
    "table_functions": table_functions,
}

docx_template.render(context)
docx_template.save(REPORT_OUTPUT_PATH)
# %%
